package barratt;

import JSci.maths.statistics.SampleStatistics;
import org.math.calculas.*;
import org.math.data.DataValue;
import org.math.data.ScalarDataValue;
import org.math.data.storage.DoubleStorage;
import org.math.data.storage.StorageType;
import org.math.factory.Combinatorics;
import org.math.factory.DataValueFactory;
import org.math.series.Series;
import org.stat.model.Model;
import org.stat.model.comparison.MSEComparable;
import org.stat.model.data.Record;
import org.stat.model.function.ModelFunction;
import org.stat.model.parameter.ParamNotEstimatedException;
import org.stat.model.parameter.Parameter;

/**
 *
 * @author Susanta Tewari
 */
public class BarrattInterferenceModel extends Model implements MSEComparable {

    double maxX = 0;
    double xBar = 0;
    double yBar = 0;

    /**
     * the default value for this scale is xBar/(yBar * 100). The scale is used
     * in the following formula x = x/xScale in {@link#compute(Record)}.
     * Optionally, this value can also be specified from the client code.
     */
    private double xScale = 0;

    public BarrattInterferenceModel(final BarrattModelData data) {

        setName("Barrat Interference Model");

        final Parameter kParam = new Parameter("k",
                                     DataValueFactory.createScalarDataValue(StorageType.DOUBLE, 0));

        setParameters(new Parameter[] { kParam });
        setModelData(data);
        setModelFunction(new BarratModelFunction());
    }

    public void setXScale(final double xScale) {
        this.xScale = xScale;
    }

    @Override
    public void runModel() {

        this.maxX = getModelData().getModelDataSummary().getMaxX()[0];
        this.xBar = getModelData().getModelDataSummary().getXBar()[0];
        this.yBar = getModelData().getModelDataSummary().getYBar()[0];

        super.runModel();
    }

    public double getMSE() {

        runModel();

        final Record[] records            = getModelData().getRecords();
        final SampleStatistics sampleStat = new SampleStatistics();

        for (final Record record : records) {

            final double observed  = record.getData()[0].getValue();
            final double estimated = record.getData()[0].getEstimatedValue();

            sampleStat.update(observed - estimated);
        }

        return sampleStat.getVariance();
    }

    /**
     * computes the following formula: y(x) = 2/3 * exp(-2x) * S_1(x)/S_k(x)
     * T_k(x) where
     * <p>
     * y : physical distance and x: recombination distance x is scaled by : x =
     * x / (xBar/yBar) (Following Dr. Arnold`s <arnold@uga.edu> advice) which is
     * the physical size on average of a map unit S_k(x) = Sum(r = 1 to
     * infinity) (2x)^r / r! * exp(-2x) * k^(r-1) T_k(x) = Sum(r = 1 to
     * infinity) r * (2x)^r / r! * [ 1 - (-1/2)^r ] * k^(r-1) and k is a
     * constant.
     */
    private class BarratModelFunction extends ModelFunction {

        double seriesEpsilon = 1.0E-10;
        int limitN           = 10;

        /** Field description */
        private final ScalarVariable varX = new ScalarVariable("x");

        /** Field description */
        private final S1Function s1Function = new S1Function();

        /** Field description */
        private final SkFunction skFunction = new SkFunction();

        /** Field description */
        private final TkFunction tkFunction = new TkFunction();
        Series s1Series                     = new Series(s1Function, seriesEpsilon, limitN);
        Series skSeries                     = new Series(skFunction, seriesEpsilon, limitN);
        Series tkSeries                     = new Series(tkFunction, seriesEpsilon, limitN);

        /** Field description */
        private double x;

        /** Field description */
        private double k;

        /**
         * The record has 2 <code>Data</code>. The first one is physical dist.
         * and the 2nd one is recom dist.
         *
         * @param record
         * @throws org.stat.exception.ParamNotEstimatedException
         */
        @Override
        public void compute(final Record record) throws ParamNotEstimatedException {

            this.k = this.model.getParameters()[0].getEstimatedValue().getValue();
            x      = record.getData(BarrattModelData.X).getValue();

            if ((xBar == 0) || (yBar == 0)) {
                throw new RuntimeException("xBar or yBar is 0");
            }

            if (xScale == 0) {


                // set the default:
                // xScale = xBar * (yBar * 100.0);
                xScale = maxX;
            }

            x /= xScale;

            varX.setValue(new ScalarDataValue(new DoubleStorage(x)));
            skFunction.setK(k);
            tkFunction.setK(k);

            final double s1SeriesVal = s1Series.evaluate().getValue();
            final double skSeriesVal = skSeries.evaluate().getValue();
            final double tkSeriesVal = tkSeries.evaluate().getValue();

            if (!s1Series.isConverged() ||!skSeries.isConverged() ||!tkSeries.isConverged()) {

                final StringBuilder strBuilder = new StringBuilder();

                strBuilder.append("x: " + x + " k: " + this.k);

                if (!s1Series.isConverged()) {

                    strBuilder.append("S1 Series didn`t converge. Iterations used: "
                                      + s1Series.getIterationUsed() + " ");
                }

                if (!skSeries.isConverged()) {

                    strBuilder.append("Sk Series didn`t converge. Iterations used: "
                                      + skSeries.getIterationUsed() + " ");
                }

                if (!tkSeries.isConverged()) {

                    strBuilder.append("Tk Series didn`t converge. Iterations used: "
                                      + tkSeries.getIterationUsed() + " ");
                }

                strBuilder.append("\n");
                logger.error(strBuilder.toString());
            }

            final double y = 2.0 / 3.0 * Math.exp(-2 * x) * s1SeriesVal / skSeriesVal * tkSeriesVal;

            record.getData()[0].setEstimatedValue(y);
        }

        /**
         * Class description
         *
         * @version        Enter version here..., 12/12/02
         * @author         Susanta Tewari
         */
        private class S1Function implements IntegerFunction {

            /** Field description */
            private final ScalarIntegerVariable varN = new ScalarIntegerVariable("n");

            public S1Function() {}

            public IntegerVariable getIntegerVariable() {
                return this.varN;
            }

            public DataValue eval() {

                final int r    = (int) varN.getValue().getValue();
                final double x = varX.getValue().getValue();
                final double result = Math.pow(2.0 * x, r) / Combinatorics.factorial(r)
                                      * Math.exp(-2 * x);

                return DataValueFactory.createScalarDataValue(StorageType.DOUBLE, result);
            }

            public Variable[] getVariables() {
                return new Variable[] { varX, varN };
            }
        }

        /**
         * Class description
         *
         * @version        Enter version here..., 12/12/02
         * @author         Susanta Tewari
         */
        private class SkFunction implements IntegerFunction {

            /** Field description */
            private final ScalarVariable varK = new ScalarVariable("k");

            /** Field description */
            private final ScalarIntegerVariable varN = new ScalarIntegerVariable("n");

            public SkFunction() {}

            public void setK(final double k) {
                varK.setValue(new ScalarDataValue(new DoubleStorage(k)));
            }

            public IntegerVariable getIntegerVariable() {
                return this.varN;
            }

            public DataValue eval() {

                final int r    = (int) varN.getValue().getValue();
                final double k = varK.getValue().getValue();
                final double x = varX.getValue().getValue();
                final double result = Math.pow(2.0 * x, r) / Combinatorics.factorial(r)
                                      * Math.exp(-2 * x) * Math.pow(k, r - 1);

                return DataValueFactory.createScalarDataValue(StorageType.DOUBLE, result);
            }

            public Variable[] getVariables() {
                return new Variable[] { varK, varX, varN };
            }
        }

        /**
         * Class description
         *
         * @version        Enter version here..., 12/12/02
         * @author         Susanta Tewari
         */
        private class TkFunction implements IntegerFunction {

            /** Field description */
            private final ScalarVariable varK = new ScalarVariable("k");

            /** Field description */
            private final ScalarIntegerVariable varN = new ScalarIntegerVariable("n");

            public TkFunction() {}

            public void setK(final double k) {
                varK.setValue(new ScalarDataValue(new DoubleStorage(k)));
            }

            public IntegerVariable getIntegerVariable() {
                return this.varN;
            }

            public DataValue eval() {

                final int r    = (int) varN.getValue().getValue();
                final double k = varK.getValue().getValue();
                final double x = varX.getValue().getValue();
                final double result = r * Math.pow(2.0 * x, r) / Combinatorics.factorial(r)
                                      * (1 - Math.pow(-0.5, r)) * Math.pow(k, r - 1);

                return DataValueFactory.createScalarDataValue(StorageType.DOUBLE, result);
            }

            public Variable[] getVariables() {
                return new Variable[] { varK, varX, varN };
            }
        }
    }
}
